import torch
import torch.nn as nn
import torch.nn.functional as F

# result["cls_output"] bert的输出
# data["bert_input"] 是所有输入
# labels 是样本标签，其中 0 表示正样本，1 表示负样本
# hyper_center 是球心的坐标

# 定义一个自定义损失函数
class CustomLoss(nn.Module):
    def __init__(self):
        super(CustomLoss, self).__init__()
    
    def forward(self, outputs, labels, hyper_center):
        mse_criterion = nn.MSELoss()
        
        # 正样本的索引
        positive_indices = (labels == 0)

        # 负样本的索引
        negative_indices = (labels == 1)

        positive_loss = 0.0
        negative_loss = 0.0

        # 正样本与球心的距离
        positive_outputs = outputs[positive_indices].squeeze()
        if(positive_outputs.shape[0] == 0):
            pass
        else:
            positive_loss = mse_criterion(positive_outputs, hyper_center)
            # positive_distances = torch.norm(positive_outputs - hyper_center, dim=1)
            # # 正样本的损失：距离的平方和
            # positive_loss = torch.sum(positive_distances ** 2)
        
        # 负样本与球心的距离
        negative_outputs = outputs[negative_indices].squeeze()
        if negative_outputs.shape[0] == 0:
            pass
        else:
            # negative_distances = torch.norm(negative_outputs - hyper_center, dim=1)
            negative_distances = mse_criterion(negative_outputs, hyper_center)
            # 负样本的损失：log(1 - exp(-distance)) 改进：-exp(-distance)
            # negative_loss = torch.sum(torch.log(1 - torch.exp(-negative_distances)))
            # negative_loss = torch.sum(-torch.exp(-negative_distances))
            # negative_loss = -torch.exp(negative_distances)
            # negative_loss = -torch.log(1 + torch.exp(-0.1 * negative_distances))
            # negative_loss = -torch.exp(-0.1 * negative_distances)
            # negative_loss = -torch.log(1 + 1.0 / negative_distances)
            # negative_loss = torch.log(1 + negative_distances)
            negative_loss = torch.log(1 + torch.exp(-negative_distances))
            # 有提升！！！
            # negative_loss =  -torch.exp(-negative_distances)

        # 整体损失：正样本损失减去负样本损失，然后求平均
        # total_loss = (positive_loss - negative_loss) / (outputs.shape[0])
        total_loss = positive_loss + negative_loss

        if(positive_outputs.shape[0] == 0):
            res_dis = negative_distances
        else:
            res_dis = positive_loss
        return total_loss, res_dis

# # 假设 hyper_center 已经定义
# hyper_center = torch.tensor([0.0, 0.0, 0.0])  # 示例球心

# # 自定义损失函数实例
# custom_loss = CustomLoss(hyper_center)

# # 生成随机样本数据
# outputs = torch.randn(10, 3)  # 假设 10 个样本，每个样本 3 维
# labels = torch.tensor([0, 0, 1, 1, 0, 0, 1, 1, 0, 1])  # 标签：0 为正样本，1 为负样本

# # 计算损失
# loss = custom_loss(outputs, labels)

# print("Loss:", loss.item())
